
from helpFunc import *
import socket,select,time,struct
from collections import deque
from base import base
import logging
from logging.handlers import RotatingFileHandler
import threading,queue
recv_q = queue.Queue()
import random

logger = logging.getLogger('my_logger')
handler = RotatingFileHandler('my_log.log', maxBytes=50*1024, backupCount=3)
formatter = logging.Formatter("%(asctime)s %(message)s",
                              "%Y-%m-%d %H:%M:%S")
handler.setFormatter(formatter)
logger.addHandler(handler)

port_g = 0
key_g = b''
one_g = (-1,-1,-1,-1)
lock = threading.Lock()
import copy  
import hashlib

class scapyRec():
    def __init__(self,n):      
        self.sock = s = socket.socket()
        add = ('127.0.0.1',n)
        s.connect(add)
        print ('scapyRec ok')   
        self.updateTime = 0
        self.updateTime2 = 0
        self.portMap = {}
        self.key = -1     
        self.rate = 0.05
        a = sniff(prn=self.aa, filter="tcp and host %s and port %s"%(serverIp,port_g) ,store=0,iface=con_iface)    
        
    def refresh(self):
        self.updateTime=time.monotonic()
        global one_g
        re = []
        for one in sorted(list(self.portMap.keys())):
            v = self.portMap[one]
            if v['updateTime']+1000<time.monotonic():
                del self.portMap[one]
                continue
            if v['updateTime']+2<time.monotonic():
                continue
            if (v['peerAck']>v['iniSeq']+3000 or(v['peerAck']>0 and v['peerAck']<v['iniSeq']-3000 and  v['peerAck']+0xffffffff>v['iniSeq']+3000)) and \
             (v['peerSeq']>v['iniAck']+3000 or(v['peerSeq']>0 and v['peerSeq']<v['iniAck']-3000 and  v['peerSeq']+0xffffffff>v['iniAck']+3000)):
                se = v['peerAck']-2000
                if se<0:
                    se += 0xffffffff    
                re.append((one,se,v['ack'],v['win']))
                if len(re)==1:
                    break
        if not re:            
            if self.key != -1:
                print ('update key',-1)           
            self.key = -1   
            lock.acquire()                         
            one_g=(self.key,-1,-1,-1)
            lock.release()            
            return

        one,s,a,w = re[0]                              
        if self.key != one:
            print ('update key',one)  
        self.key = one              
        lock.acquire()                         
        one_g=(self.key,s,a,w)
        lock.release()                             

    def aa(self,p):   
        tt = time.monotonic()
        t = p[TCP]
        p = t.payload.original
        dk = hashlib.pbkdf2_hmac('md5', p[:-16], con_key, 2)      
        isMy = p[-16:]== dk        
        if isMy and p[:6]==key_g and t.sport == port_g:    
            if recv_q.qsize()>50:
                return
            recv_q.put(p[6:-16])
            self.sock.sendall(b'1')
            return
        if isMy:
            return       
        xx = str(t.flags)
        if 'F' in xx or 'R' in xx:
            if t.dport!=port_g  and t.dport in self.portMap:
                del self.portMap[t.dport]
                print ('del key',t.dport)
            if t.sport!=port_g  and t.sport in self.portMap:
                del self.portMap[t.sport] 
                print ('del key',t.sport)
            return self.refresh()
        
        if 'S' in xx and t.dport==port_g:
            self.portMap[t.sport] = {'iniSeq':t.seq,'iniAck':-1,'seq':-1,'ack':-1,'peerAck':-1,'peerSeq':-1,'win':-1,
                                     'updateTime':tt}
        
        if 'S' in xx and 'A' in xx and t.sport==port_g and t.dport in self.portMap:
            self.portMap[t.dport]['iniAck'] = t.seq
            return self.refresh()  
        
        if t.dport==port_g and  t.sport in self.portMap:
            self.portMap[t.sport]['seq'] = t.seq
            self.portMap[t.sport]['ack'] = t.ack
            self.portMap[t.sport]['win'] = t.window
            self.portMap[t.sport]['updateTime'] = tt
            if tt>self.updateTime+self.rate:
                return self.refresh()
        
        if t.sport==port_g and  t.dport in self.portMap:
            self.portMap[t.dport]['peerAck'] = t.ack
            self.portMap[t.dport]['peerSeq'] = t.seq
            self.portMap[t.dport]['updateTime'] = tt
            if tt>self.updateTime+self.rate:
                return self.refresh()              
                 
class cli(base):
    def __init__(self,p1,p2,onePakcLen,packEffectTime,timeoutInter,port,k):
        super().__init__(p1,p2,onePakcLen,packEffectTime,timeoutInter)
        global port_g,key_g
        key_g = k
        port_g = port     
        self.sockMap = {}
        self.staTime = time.monotonic()
        self.staGot = self.staOut = 0
        self.udpCo = 0
        self.lastCheckTimeout = 0
        self.newSockThisSecond = 0
        self.recvTime = 0
        self.getNumPerPack = con_getNum
        self.peerMaxPos = 0    
        self.recvMap = {}
        self.rRaw =self.wRaw = 0
        self.timeoutTime = timeoutTime
        self.peerMaxGotLen = self.peerMaxSendLen = 0
        self.nullSend = 0        
        self.staSend = 0    
        self.newSockLimit = maxNewSockPerSecond
        self.controlSock = self.iniControlSock()
        self.staRecv = 0
        self.packNotGot = 0
        self.serverTotalSendLen = serverSenLen
        self.serverOneLen = serverPackLen
        self.totalSend = 0
        self.dSock = None
        self.iniDsock()
        self.scapySock = conf.L2socket(iface=con_iface)    
        self.notInMapNum = 0
        self.notSendNum = 0
        self.realOut = 0
        import plink
        def ff():
            plink.main()
        self.pl = x = threading.Thread(target=ff)
        x.setDaemon(True)
        x.start() 
        
    def send(self,d):
        lock.acquire()
        sport,seq,ack,win = copy.deepcopy(one_g)
        lock.release()
        if sport<0:
            self.notSendNum += 1
            return
        ip = IP(dst=serverIp)
        st = key_g+d
        dk = hashlib.pbkdf2_hmac('md5', st, con_key, 2)        
        tcp = Ether() / ip / TCP(sport=sport, dport=port_g, flags="PA", seq=seq, ack=ack,window=win)/(st+dk)
        self.scapySock.send(tcp)
        
    def iniDsock(self):
        s = socket.socket()
        n = find_free_port(socket.SOCK_STREAM)
        add = ('127.0.0.1',n)
        s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
        s.bind(add)
        s.listen(1)
        def func():
            d = scapyRec(n)
        self.t = t = threading.Thread(target=func)
        t.setDaemon(True)
        t.start()     
        self.dSock,b = s.accept()
        s.close()        
        
    def iniControlSock(self):
        s = socket.socket()
        add = ('127.0.0.1',con_controlPort)
        s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)
        s.bind(add)
        s.listen(1)
        return s
    
    def getId(self):
        self.udpCo+=1
        if self.udpCo==60000:
            self.udpCo=0
        return struct.pack('H',self.udpCo)
    
    def dealSta(self,ss):
        l = []
        mi = struct.unpack('Q',ss[:8])[0]
        if mi >self.peerPos:
            le = 0
            for i in range(self.peerPos,mi):
                le += len(self.sendMap[i]['data'])
                del self.sendMap[i]
            self.peerPos = mi
            self.tellAddBuf(le)            
        ma = struct.unpack('Q',ss[8:16])[0]
        d = ss[16:]
        co = -1
        pos = 0
        b = 0
        for i in range(mi,ma+1):
            co += 1
            if co==0:
                b = d[pos]
                pos += 1
                if b&128:
                    l.append(i)
                b <<= 1
            else:
                if b&128:
                    l.append(i)
                b <<= 1
                if co==7:
                    co = -1
        for one in l:
            if one in self.sendMap:
                self.sendMap[one]['peerGot'] = True

    def dealRecvPack(self,r):
        w = structWrapper(r)
        id = w.readArr(2)        
        if id not in self.sockMap:
            self.notInMapNum += 1
            return        
        askList = self.sockMap[id]['askList'] 
        dataList = self.sockMap[id]['dataList']         
        if dataList and dataList[0] in self.sendMap:
            self.sendMap[dataList[0]]['peerGot'] = True                     
        self.lastRecvTime = time.monotonic()
        self.staGot+=1  
        self.rRaw += len(r)
        if len(r)>self.maxGotLen:
            self.maxGotLen = len(r)                     
        tt = time.monotonic()-self.sockMap[id]['createTime']
        if tt>self.recvTime:
            self.recvTime = tt
        mm = w.readQ()
        if mm >self.peerMaxPos:
            self.peerMaxPos=mm
        pmg = w.readWord()
        if pmg>self.peerMaxGotLen:
            self.peerMaxGotLen=pmg
        pmg = w.readWord()
        if pmg>self.peerMaxSendLen:
            self.peerMaxSendLen = pmg
        sendNum = w.readByte()
        for i in range(len(askList)):
            nn = askList[i]
            if i<sendNum:
                le = w.readWord()
                st = w.readArr(le)                    
            if nn not in self.recvMap or self.recvMap[nn]['got']:
                continue
            if i<sendNum and st:
                self.recvMap[nn]['got'] = True
                self.recvMap[nn]['data'] = st
                self.checkRecvMap()                
            else:
                self.recvMap[nn]['askTime'] -= 1
                self.packNotGot += 1
        hasSta = w.readByte()
        if hasSta == 1:
            self.staRecv += 1
            le = w.readWord()
            sta = w.readArr(le)
            self.dealSta(sta)
        addPackNum = w.readByte()
        for i in range(addPackNum):
            nn = w.readQ()
            ll = w.readWord()
            st = w.readArr(ll)
            if nn<self.recvPos:
                continue
            self.recvMap[nn] = {'got':True,'askTime':0,'data':st}
        self.checkRecvMap() 
        del self.sockMap[id]
        
    def doPeriodWork(self):    
        if time.monotonic()-self.lastRecvTime>120:
            import os
            os._exit(0)          
        p1,p2 = int(self.rRaw/1024),int(self.wRaw/1024)
        p3,p4 = int(self.rNet/1024),int(self.wNet/1024)
        tt = int(self.recvTime*100)/100
        self.timeoutTime = self.recvTime+con_timeoutPlus
        if self.recvTime==0 or self.recvTime>timeoutTime:
            self.timeoutTime = timeoutTime
        if self.realOut==0:
            self.gotRate = 1
        else:
            self.gotRate = self.staGot/self.totalSend
        self.maxAskTime = self.calAskTime()
        st = '%s %s || %s %s %s || %s %s %s %s || %s %s %s %s <null sta loss> %s %s %s <q,map,rej> %s %s %s || %s %s %s %s'\
        %(len(self.sockMap),tt,self.staGot,self.realOut, self.staOut,p1,p2,p3,p4,self.maxSenLen,self.peerMaxGotLen,self.maxGotLen,\
          self.peerMaxSendLen,self.nullSend,self.staRecv,self.packNotGot,recv_q.qsize(),self.notInMapNum,self.notSendNum,\
          self.sendPos,self.peerPos,self.peerMaxPos,self.recvPos)
        logger.warning(st)
        print (st)
        self.staGot = self.staOut = self.newSockThisSecond= self.recvTime=self.staSend=self.notInMapNum=0
        self.rRaw =self.wRaw = 0
        self.rNet = self.wNet = 0
        self.nullSend = self.packNotGot = self.totalSend = self.notSendNum =self.realOut=0
        self.peerMaxGotLen = self.peerMaxSendLen = self.maxGotLen=self.maxSenLen=self.staRecv=0         
        
    def sendPacks(self):
        passTime = time.monotonic()-self.staTime
        if passTime >1:
            passTime = 1
        n = passTime*self.newSockLimit-self.newSockThisSecond
        n = int(n)
        if n<0:
            n=0    
        if self.gotRate==0:
            minPackNum = 20000
            maxSend =  20000
        else:            
            minPackNum = passTime*(serverStaPerSecond/self.gotRate)
            maxSend =  passTime*(con_maxPackRecv/self.gotRate)
        maxSend =  int(maxSend-self.totalSend)
        if n>maxSend:
            n = maxSend            
        staSendNum = int(minPackNum-self.staSend)        
        if staSendNum<0:
            staSendNum=0
        if staSendNum>n:
            staSendNum=n
        self.staSend += staSendNum
        askList = self.getAskList(n,staSendNum)
        dataList = self.getDataToSend(n)                
        for i in range(n):
            al = askList[i:i+1]
            dl = dataList[i:i+1]
            if staSendNum>0:
                staSendNum -= 1
                hasSta = 1
            else:
                hasSta = 0                
                if not al and not dl:
                    continue
            self.newSockThisSecond += 1  
            self.totalSend += 1
            self.packMsg_send(al,dl,hasSta)
            if not al and not dl:
                self.nullSend += 1     
                
    def dealControl(self):
        try:
            a,b = self.controlSock.accept()
            st = b''
            while len(st)!=2:
                st += a.recv(2)
            self.newSockLimit = struct.unpack('H',st)[0]
            print ('newSockLimit changed',self.newSockLimit)
            a.close()
        except Exception as e:
            print (e)
    def work(self):       
        rl,wl = self.rig()
        rl.append(self.controlSock)
        rl.append(self.dSock)    
        rl,wl,el = select.select(rl,wl,[],self.calWaitTime()/1000)
        self.dealFd(rl,wl)
        if self.controlSock in rl:
            self.dealControl()
        if self.dSock in rl:
            r = self.dSock.recv(tcpReadLen)
            for i in range(len(r)):
                d = recv_q.get()
                self.dealRecvPack(d)            
        self.reFreshSendMap()
        if time.monotonic()-self.lastCheckTimeout>con_timeoutInternal:
            self.lastCheckTimeout = time.monotonic()
            self.dealTimeout()
        self.sendPacks()                    
        if time.monotonic()-self.staTime>1:
            self.staTime = time.monotonic()
            self.doPeriodWork()
     
    def checkRecvMap(self):
        while self.recvPos in self.recvMap and self.recvMap[self.recvPos]['got']:
            self.wrietBuf += self.recvMap[self.recvPos]['data']
            del self.recvMap[self.recvPos]
            self.recvPos += 1        
            
    def packMsg_send(self,askList,dataList,hasSta):
        if askList:
            askList = sorted(askList[0])
        tt = time.monotonic()
        id = self.getId()     
        w = structWrapper()    
        w.writeArr(id)
        w.writeQ(self.recvPos)
        w.writeFloat(self.gotRate)
        if not dataList:
            w.writeByte(0)
        else:
            w.writeByte(1)
            w.writeQ(dataList[0])
            w.writeWord(len(self.sendMap[dataList[0]]['data']))
            w.writeArr(self.sendMap[dataList[0]]['data'])
        w.writeByte(hasSta)
        le = len(askList)
        w.writeByte(le)
        if le!=0:
            w.writeQ(askList[0])
            for i in range(le-1):
                w.writeDWord(askList[i+1]-askList[0])
        if id in self.sockMap:
            print ('error')
            os._exit(0)
        self.sockMap[id]={'askList':askList, 'dataList':dataList,'createTime':tt,'dead':False}
        self.send(w.data)
        self.wRaw += len(w.data)
        if len(w.data)>self.maxSenLen:
            self.maxSenLen = len(w.data)                 
    
    def getAskList(self,n,staPackNum):
        if n == 0:
            return []
        staLen = 17+(self.sendPos-self.peerPos)//8
        sendNum = (self.serverTotalSendLen-20-staLen)//(self.serverOneLen+2)
        if sendNum<=0:
            sendNum=0
        l = []
        tempList = []
        thisLoopSendTime = 0
        while True:
            thisLoopSendTime+=1
            if thisLoopSendTime > self.maxAskTime:
                if tempList and len(l)<n:
                    l.append(tempList)
                return l
            for i in range(self.recvPos,self.peerMaxPos):
                if staPackNum>0 and len(tempList)>=sendNum:
                    staPackNum-=1
                    l.append(tempList)
                    tempList = []
                    if len(l)==n:
                        return l                             
                if i not in self.recvMap:
                    self.recvMap[i] = {'got':False,'askTime':0,'data':b''}
                if not self.recvMap[i]['got'] and self.recvMap[i]['askTime']<thisLoopSendTime:
                    if i in tempList:
                        l.append(tempList)
                        tempList = []
                        if len(l)==n:
                            return l                        
                    self.recvMap[i]['askTime']+=1
                    tempList.append(i)
                    if len(tempList)==self.getNumPerPack:
                        l.append(tempList)
                        tempList = []
                        if len(l)==n:
                            return l    
                        
    def getDataToSend(self,n):
        if n == 0:
            return []        
        l = []
        thisLoopSendTime = 0
        while True:
            thisLoopSendTime+=1
            if thisLoopSendTime > self.maxAskTime:
                return l
            for i in range(self.peerPos,self.sendPos):
                if not self.sendMap[i]['peerGot'] and self.sendMap[i]['sendTime']<thisLoopSendTime:
                    self.sendMap[i]['sendTime']+=1
                    l.append(i)
                    if len(l)==n:
                        return l        
        
    def addToDeadSock(self,k,doDel):  
        if doDel and self.sockMap[k]['dead'] :
            del self.sockMap[k]
            return
        askList = self.sockMap[k]['askList']                   
        dataList = self.sockMap[k]['dataList']    
        for i in askList:
            if i in self.recvMap and not self.recvMap[i]['got']:
                self.recvMap[i]['askTime'] -= 1
        if dataList and dataList[0] in self.sendMap and not self.sendMap[dataList[0]]['peerGot']:
            self.sendMap[dataList[0]]['sendTime'] -= 1                
        if doDel:
            del self.sockMap[k]
        else:
            self.sockMap[k]['dead']  = True

    def dealTimeout(self):
        tt = time.monotonic()
        for k in list(self.sockMap.keys()):
            if self.sockMap[k]['createTime']+self.timeoutTime<tt and not self.sockMap[k]['dead']:
                self.addToDeadSock(k,False)
                self.staOut += 1  
            if self.sockMap[k]['createTime']+2<tt:                
                self.addToDeadSock(k,True)
                self.realOut += 1
                
if __name__ == '__main__':
    import threading
    l = findNPorts(2,socket.SOCK_STREAM)
    def ww():
        a = cli(l[0],l[1],cliPackLen,
                 con_packEffectTime,con_timeoutInternal, 22,b'sdfeff')
        while True:
            a.work()
    t = threading.Thread(target=ww)
    t.setDaemon(True)
    t.start()
    from testconn import testconn
    xx = testconn(l[0],l[1])


                                     

